import random

import cv2
import numpy as np
import matplotlib.pyplot as plt


def filter_out_red(img_input):
    if img_input is not None:
        mask = np.zeros(img_input.shape[:-1]).astype(np.uint8)
        for row in range(img_input.shape[0]):
            for col in range(img_input.shape[1]):
                # if 0 <= hsv[row, col, 0] <= 10 or 156 <= hsv[row, col, 0] <= 180:
                #     if 43 <= hsv[row, col, 1] <= 255 and 46 <= hsv[row, col, 2] <= 255:
                if pixel_is_red(img_input[row, col]):
                    mask[row, col] = 255
        return cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR), cv2.bitwise_and(img_input, img_input, mask=mask)


def filter_out_black(img_input):
    if img_input is not None:
        hsv = cv2.cvtColor(img_input, cv2.COLOR_BGR2HSV)
        # img_out = (np.ones(img_input.shape) * 255).astype(np.uint8)
        mask = np.zeros(img_input.shape[:-1]).astype(np.uint8)
        for row in range(img_input.shape[0]):
            for col in range(img_input.shape[1]):
                if hsv[row, col, 2] <= 100:
                    # img_out[row, col] = img_input[row, col]
                    mask[row, col] = 255
        return cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR), cv2.bitwise_and(img_input, img_input, mask=mask)


def filter_out_gray(img_input):
    if img_input is not None:
        hsv = cv2.cvtColor(img_input, cv2.COLOR_BGR2HSV)
        img_out = (np.ones(img_input.shape) * 255).astype(np.uint8)
        for row in range(img_input.shape[0]):
            for col in range(img_input.shape[1]):
                if hsv[row, col, 1] <= 43 and 46 <= hsv[row, col, 2] <= 220:
                    img_out[row, col] = img_input[row, col]
        return img_out


def filter_out_orange(img_input):
    if img_input is not None:
        hsv = cv2.cvtColor(img_input, cv2.COLOR_BGR2HSV)
        lower = np.array([11, 43, 46])
        upper = np.array([25, 255, 255])
        mask = cv2.inRange(hsv, lower, upper)
        return mask, cv2.bitwise_and(img_input, img_input, mask=mask)


def filter_out_yellow(img_input):
    if img_input is not None:
        hsv = cv2.cvtColor(img_input, cv2.COLOR_BGR2HSV)
        lower = np.array([26, 43, 46])
        upper = np.array([34, 255, 255])
        mask = cv2.inRange(hsv, lower, upper)
        return mask, cv2.bitwise_and(img_input, img_input, mask=mask)


def pixel_is_red(bgr):
    hsv = cv2.cvtColor(bgr.reshape(1, 1, 3), cv2.COLOR_BGR2HSV)
    if 0 <= hsv[0, 0, 0] <= 10 or 156 <= hsv[0, 0, 0] <= 180:
        if 43 <= hsv[0, 0, 1] <= 255 and 46 <= hsv[0, 0, 2] <= 255:
            return True
    return False


def pixel_is_black(bgr):
    hsv = cv2.cvtColor(bgr.reshape(1, 1, 3), cv2.COLOR_BGR2HSV)
    if hsv[0, 0, 2] <= 46:
        return True
    return False


def get_red_pixel(from_list=False):
    if from_list:
        rgb_list = [[156, 50, 76], [102, 31, 35], [102, 28, 53], [107, 21, 34], [134, 35, 44],
                    [185, 118, 124], [205, 94, 125], [113, 78, 84], [110, 70, 81], [70, 16, 14],
                    [107, 48, 44], [53, 7, 7]]
        rgb = random.choice(rgb_list)
        rgb = np.array(rgb, dtype=np.uint8).reshape(1, 1, 3)
        temp = rgb[0, 0, 0]
        rgb[0, 0, 0] = rgb[0, 0, 2]
        rgb[0, 0, 2] = temp
        bgr = rgb
        return bgr
    else:
        hsv = np.zeros((1, 1, 3), dtype=np.uint8)

        if np.random.rand() < 0.5:
            hsv[0, 0, 0] = np.random.randint(0, 10)
        else:
            hsv[0, 0, 0] = np.random.randint(156, 180)
        hsv[0, 0, 1] = np.random.randint(43, 255)
        hsv[0, 0, 2] = np.random.randint(46, 255)
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        return bgr


def get_red_pixel2():
    hsv = np.zeros((1, 1, 3), dtype=np.uint8)

    if np.random.rand() < 0.5:
        hsv[0, 0, 0] = np.random.randint(0, 10)
    else:
        hsv[0, 0, 0] = np.random.randint(156, 180)
    hsv[0, 0, 1] = np.random.randint(200, 255)
    hsv[0, 0, 2] = np.random.randint(46, 255)
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    return bgr


def get_black_pixel():
    hsv = np.zeros((1, 1, 3), dtype=np.uint8)
    hsv[0, 0, 0] = np.random.randint(0, 180)
    hsv[0, 0, 1] = np.random.randint(0, 255)
    hsv[0, 0, 2] = np.random.randint(0, 46)
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    return bgr


def get_number_color(img, black_rate):
    area_number = img[int(img.shape[0] * 3 / 8): int(img.shape[0] * 5 / 8),
                  int(img.shape[1] * 3 / 8): int(img.shape[1] * 5 / 8), :]
    hsv = cv2.cvtColor(area_number, cv2.COLOR_BGR2HSV)
    h, s, v, cnt = 0, 0, 0, 0
    th, img_b = cv2.threshold(hsv[:, :, 2], 0, 255, cv2.THRESH_OTSU)
    for row in range(hsv.shape[0]):
        for col in range(hsv.shape[1]):
            if hsv[row, col, 2] < th:
                h += hsv[row, col, 0]
                s += hsv[row, col, 1]
                v += hsv[row, col, 2]
                cnt += 1
    hsv_mean = np.zeros((1, 1, 3))
    hsv_mean[0, 0, 0] = h / cnt
    hsv_mean[0, 0, 1] = s / cnt
    hsv_mean[0, 0, 2] = black_rate + v / cnt
    bgr = cv2.cvtColor(hsv_mean.astype(np.uint8), cv2.COLOR_HSV2BGR)
    return int(bgr[0, 0, 0]), int(bgr[0, 0, 1]), int(bgr[0, 0, 2])


def red2black_circle(img, color=get_black_pixel()):
    img_cp = img.copy()
    h, w, c = img.shape
    cnt = 0
    for row in range(h):
        for col in range(w):
            if (w / 6 + h / 6) ** 2 < (row - h / 2) ** 2 + (col - w / 2) ** 2 < (
                    (w + h) * 11 / 48) ** 2 and pixel_is_red(img[row, col]):
                img[row, col] = color
                cnt += 1
    if cnt > 5 * (w + h) / 3:
        return img
    else:
        return img_cp


def black2red_circle(img, color=get_red_pixel()):
    img_cp = img.copy()
    h, w, c = img.shape
    cnt = 0
    for row in range(h):
        for col in range(w):
            if (w / 6 + h / 6) ** 2 < (row - h / 2) ** 2 + (col - w / 2) ** 2 < (
                    (w + h) * 11 / 48) ** 2 and pixel_is_black(img[row, col]):
                img[row, col] = color
                cnt += 1
    if cnt > 5 * (w + h) / 3:
        return img
    else:
        return img_cp


def limit0_255(x):
    if x < 0:
        x = 0
    if x > 255:
        x = 255
    return x


def myYUV2BGR(yuv):
    """
    测试C语言写的YUV转RGB是否正确
    """
    bgr = np.zeros(yuv.shape, dtype=np.uint8)
    for row in range(yuv.shape[0]):
        for col in range(yuv.shape[1]):
            y = yuv[row, col, 0]
            u = yuv[row, col, 1] - 128
            v = yuv[row, col, 2] - 128
            r = limit0_255(y + 1.402 * v)
            g = limit0_255(y - 0.344 * u - 0.714 * v)
            b = limit0_255(y + 1.772 * u)
            bgr[row, col] = b, g, r
    return bgr


def myBGR2HSV(bgr):
    hsv = np.zeros(bgr.shape, dtype=np.uint8)
    for row in range(bgr.shape[0]):
        for col in range(bgr.shape[1]):
            b, g, r = bgr[row, col] / 255
            cmax = max(b, g, r)
            cmin = min(b, g, r)
            c = cmax - cmin
            v = cmax
            if c == 0:
                h = 0
            elif r == v:
                h = 60 * (g - b) / c
            elif g == cmax:
                h = 60 * (2 + (b - r) / c)
            elif b == cmax:
                h = 60 * (4 + (r - g) / c)
            if h < 0:
                h += 360
            if v == 0:
                s = 0
            else:
                s = c / v
            h /= 2
            s *= 255
            v *= 255
            hsv[row, col] = [h, s, v]
    return hsv


def draw_line(img, line_nums, color, thickness):
    for idx in range(line_nums):
        p1 = [random.randint(int(idx * (img.shape[0] + img.shape[1]) / line_nums),
                             int((idx + 1) * (img.shape[0] + img.shape[1]) / line_nums - 1)), 0]
        p2 = [0, random.randint(0, img.shape[0] - 1)]
        p2[0] = p1[1] + p1[0] - p2[1]
        p3 = [0, random.randint(0, img.shape[0] - 1)]
        p3[0] = p1[1] + p1[0] - p3[1]
        p4 = [0, 0]
        p4[1] = p1[1] + p1[0]
        cv2.line(img, p4, p1, color, thickness)


def get_different_color(input_color, dif):
    a = [0, 0, 0]
    for i in range(3):
        if input_color[i] - dif > 0:
            a[i] = random.randint(0, input_color[i] - dif)
        else:
            a[i] = random.randint(input_color[i] - dif, 255)
    return a[0], a[1], a[2]


def motion_blur(img, degree=10, angle=20):
    image = img.copy()
    # 这里生成任意角度的运动模糊kernel的矩阵， degree越大，模糊程度越高
    M = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
    motion_blur_kernel = np.diag(np.ones(degree))
    motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (degree, degree))

    motion_blur_kernel = motion_blur_kernel / degree
    blurred = cv2.filter2D(image, -1, motion_blur_kernel)
    # convert to uint8
    cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
    blurred = np.array(blurred, dtype=np.uint8)
    return blurred


def get_vstack_src_maskred_red_img(img):
    mask_red = filter_out_red(img)
    return np.vstack((img, mask_red[0], mask_red[1]))


if __name__ == '__main__':
    file_path = r"H:\img\tsr\00006_3009_0.587182_0.991992_0.970827_00-46-27_1381-289-46-44_20-6.bmp"
    img = cv2.imread(file_path)
    mask, img_f = filter_out_yellow(img)
    img_show = np.hstack((img, cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)))
    # for row in range(img.shape[0]):
    #     for col in range(img.shape[1]):
    #         if 20 <= img[row, col] <= 80:
    #             img[row, col] = 255
    #             img[row, col].astype(np.uint8)

    # hist_gray = cv2.calcHist([img], [0], None, [256], [0, 256])
    # plt.plot(hist_gray)
    cv2.imshow('img_show', img_show)
    # plt.show()
    cv2.waitKey(0)
